print("Initializing...")

# pytorch
import torch
import torchvision.transforms as transforms

# 网络结构
from models.net import Net

# WebUI
import gradio as gr

# 图像处理
from PIL import Image

# 命令行参数获取
import utils.parameters

# 读取命令行参数
params_parser = utils.parameters.get_run_args()

print("Loading model...")

# 决定运行设备
device = "cuda:0" if torch.cuda.is_available() else "cpu"
# 实例化网络
model = Net()
# 加载模型参数
model.load_state_dict(
    torch.load(
        "./checkpoints/" + params_parser.model_name,
        weights_only=True,
        map_location=torch.device(device),
    )
)
# 设置模型加载到设备
model = model.to(device)
model.eval()  # 切换到推理模式

print("Model loaded.")


# 定义图像预处理函数
def preprocess_image(image):
    # 将图像转换为RGB图像
    image = image.convert("RGB")
    # 缩放到32x32像素
    image = image.resize((32, 32))
    # 转换为张量
    image = transforms.ToTensor()(image)
    # 把第一个批次添加到设备
    image = image.unsqueeze(0).to(device)
    return image


# 定义预测函数
def predict_digit(image):
    # 预处理图像
    image = preprocess_image(image)
    # 前向传播
    with torch.no_grad():
        output = model(image)
    # 获取预测结果
    prediction = torch.argmax(output, dim=1).item()
    label_str = [
        "airplane",
        "automobile",
        "bird",
        "cat",
        "deer",
        "dog",
        "frog",
        "horse",
        "ship",
        "truck",
    ]
    return label_str[prediction]


print("Loading Gradio...")

# 创建Gradio界面
with gr.Blocks() as demo:
    with gr.Row():
        # 创建图片上传器
        image_input = gr.Image(label="Upload the image", type="pil")
        # 创建输出文本框
        output_text = gr.Textbox(label="Prediction")
    # 创建按钮
    predict_button = gr.Button("Predict")
    # 设置按钮点击事件
    predict_button.click(fn=predict_digit, inputs=image_input, outputs=output_text)

# 启动Web服务
demo.launch(share=params_parser.share)
